package org.apache.storm.kafka;

import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.storm.generated.StormTopology;
import org.apache.storm.kafka.bolt.LogBolt;
import org.apache.storm.kafka.bolt.SimpleShellBolt;
import org.apache.storm.kafka.spout.*;
import org.apache.storm.topology.TopologyBuilder;
import org.apache.storm.tuple.Fields;
import org.apache.storm.tuple.Values;

import static org.apache.storm.kafka.spout.KafkaSpoutConfig.FirstPollOffsetStrategy.EARLIEST;

/**
 * Created by peibin on 2017/11/22.
 */
public class ConsumerTopology {
    private static final String STREAM = "stream";

    private static StormTopology newConsumerTopology(KafkaSpoutConfig<String, String> spoutConfig) {
        final TopologyBuilder tp = new TopologyBuilder();
        tp.setSpout("kafka_spout", new org.apache.storm.kafka.spout.KafkaSpout<>(spoutConfig), 2);
        tp.setBolt("shell_bolt", new SimpleShellBolt(), 4)
                .fieldsGrouping("kafka_spout", STREAM, new Fields("prefix"));
        tp.setBolt("log_bolt", new LogBolt(),4).shuffleGrouping("shell_bolt");
        return tp.createTopology();
    }

    protected static KafkaSpoutConfig<String, String> getKafkaSpoutConfig(String bootstrapServers, String topic) {
        ByTopicRecordTranslator<String, String> trans = new ByTopicRecordTranslator<>(
                (r) -> new Values(r.topic(), r.partition(), r.offset(), r.key(), r.value().substring(0, 1), r.value()),
                new Fields("topic", "partition", "offset", "key", "prefix", "value"), STREAM);
        return KafkaSpoutConfig.builder(bootstrapServers, new String[]{topic})
                .setProp(ConsumerConfig.GROUP_ID_CONFIG, "Main")
                .setRetry(getRetryService())
                .setRecordTranslator(trans)
                .setOffsetCommitPeriodMs(10_000)
                .setFirstPollOffsetStrategy(EARLIEST)
                .setMaxUncommittedOffsets(250)
                .build();
    }

    // 重试策略
    protected static KafkaSpoutRetryService getRetryService() {
        return new KafkaSpoutRetryExponentialBackoff(KafkaSpoutRetryExponentialBackoff.TimeInterval.microSeconds(500),
                KafkaSpoutRetryExponentialBackoff.TimeInterval.milliSeconds(2), Integer.MAX_VALUE, KafkaSpoutRetryExponentialBackoff.TimeInterval.seconds(10));
    }

    public static StormTopology newConsumerTopology(String brokerUrl, String topic) {
        return newConsumerTopology(getKafkaSpoutConfig(brokerUrl, topic));
    }
}
